import PySide6.QtGui
from PySide6.QtWidgets import QTextEdit
from PySide6.QtCore import Qt


class CodeTextEdit(QTextEdit):

    def __init__(self, parent, PySide6_QtWidgets_QWidget=None, *args, **kwargs):
        super().__init__(parent, PySide6_QtWidgets_QWidget, *args, **kwargs)

    def keyPressEvent(self, e: PySide6.QtGui.QKeyEvent) -> None:
        if e.key() == Qt.Key.Key_Tab:
            self.textCursor().insertText(' ' * 4)
        elif e.key() == Qt.Key.Key_Return:
            text, p = self.toPlainText(), self.textCursor().position()
            before_text = text[:p]
            insert_space = '\n' + '    ' * (before_text.count('{') - before_text.count('}'))
            self.textCursor().insertText(insert_space)
        else:
            super().keyPressEvent(e)
